#!/usr/bin/env python
#coding:utf-8
#Code by Mango

from geColor import *
from sqlite3 import *
from os import listdir,system,popen
from random import choice
from geLoad import load


gedb = 'cache/ge.cache' #数据库路径
plugp = 'plugins/'  #插件主目录
#p_cli = 'client'    # 客户端目录
#p_ser = 'server'    # 服务端目录
p_mul = 'multi'  # 辅助插件目录
pyge = 'PyGE'

class cache:
    def start(self):  # 数据库初始化函数
        color.cprint("[*] Start PyGE ...",GREEN)  # 绿色显示程序开始运行
		# 数据库初始化相关
        color.cprint("[*] SQL Ready ...",GREEN)
        self.runsql("create table if not exists pyge(id integer primary key,type text,path text)") # 新建一个表id,type,path三个字段
        # 表server , 字段id,ip,cpu,mem,memfree,swap,swapfree,disk,diskfree,day,uptime
        self.runsql("create table if not exists server(id integer primary key,ip text,cpu text,mem text,memfree text,swap text,swapfree text,disk text,diskfree text,day text,uptime text)")
        self.runsql("delete from pyge")   # 清空表的内容,以便更新数据
        self.runsql("delete from server")   # 清空表的内容,以便更新数据
        # 从ip列表读取ip初始化数据库
        f = open('cache/client.txt')  #打开主机列表文件
        while True:
            line = f.readline().split()  #读取文件,把每一行用空格分成一个列表
            if len(line) == 0:break    #读到空之后跳出循环 
            #print line[0]
            self.runsql('insert into server(ip,cpu,mem,memfree,swap,swapfree,disk,diskfree,day,uptime) values("%s","NULL","NULL","NULL","NULL","NULL","NULL","NULL","NULL","NULL")' % line[0])
        f.close()
		# 更新插件数量
        #self.inscache(self.getplus(p_cli),p_cli)
        #self.inscache(self.getplus(p_ser),p_ser)
        self.inscache(self.getplus(p_mul),p_mul)
        color.cprint("[*] SQL Ready Success ",GREEN)
        # 启动相关服务
        color.cprint("[*] Server Start ... ",GREEN)
        '''
        system("python m_server.py &")
        '''
        self.banner()

    def inscache(self,c,p):  # 插入数据函数,这个只是作为缓存插件的
        for tmp in c:
            tmp = tmp[:len(tmp)-3]
            self.runsql('insert into pyge(type,path) values("%s","%s/%s")'%(p,p,tmp))

    def runsql(self,sql):    # sql执行函数
        conn = connect(gedb)  # 连接sqldb,没有则新建一个
        conn.execute(sql)    # 执行语句
        conn.commit()    # 执行
        conn.close()    # 关闭连接

    def getplus(self,path):  # 获取路径函数
        return listdir(plugp+path)  # os.listdir 列出目录下所有文件和目录
		
    def sql_all(self,sql):  # 获取sqldb数据的函数
        conn=connect(gedb)    # 连接sqldb,没有则新建一个
        cur=conn.cursor()    #创建一个游标
        cur.execute(sql)     # 执行语句
        tmp=cur.fetchall()    #获取数据
        cur.close()
        conn.close()
        return tmp    #返回数据

    def search(self,sear):
        sql='select * from pyge where path like "%'+sear+'%"'
        result=self.sql_all(sql)
        msg="SEARCH '%s'"%sear
        color.cprint(msg,YELLOW)
        color.cprint("="*len(msg),GREY)
        self.listpyge(result)

    def search_host(self, sear):    #主机搜索
        sql = 'select * from server where ip like "%'+sear+'%"'
        result = self.sql_all(sql)
        msg="SEARCH '%s'" % sear
        color.cprint(msg, YELLOW)
        color.cprint("="*len(msg),GREY)
        self.show_ser(result)

    def listpyge(self,result):
        color.cprint("%5s %-60s %-7s"%("ID","PATH","TYPE"),YELLOW)
        color.cprint("%5s %-60s %-7s"%("-"*5,"-"*60,"-"*7),GREY)
        for res in result:
            rid=res[0]
            rty=res[1]
            rpa=res[2]
            if len(rpa)>70:
                rpa=rpa[:68]+".."
            color.cprint("%5s %-60s %-7s"%(rid,rpa,rty),CYAN)
        color.cprint("="*74,GREY)
        color.cprint("COUNT [%s] RESULTS (*^_^*)"%len(result),GREEN)

    def show_host(self, result):   # 显示所有主机
        for res in result:
            #print res,type(res)
            a = res[0]
            color.cprint("%5s " % a,CYAN)
        color.cprint("=" * 12,GREY)
        color.cprint("COUNT [%s] HOST " % len(result),GREEN)

    def show_ser(self, result):  # 显示所有主机状态
        color.cprint("%-3s %-15s %-5s %-8s %-8s %-8s %-8s %-8s"%("ID","Host","CPU","MEMFree","SwapFree","DiskFree","Day","UPTime"),YELLOW)
        color.cprint("-"*72,GREY)
        for res in result:
            a = res[0]
            b = res[1]
            c = res[2]
            d = res[3]
            e = res[4]
            f = res[5]
            g = res[6]
            h = res[7]
            color.cprint("%-3s %-15s %-5s %-8s %-8s %-8s %-8s %-8s" % (a,b,c,d,e,f,g,h),CYAN)
            #print res,type(res)
        color.cprint("=" * 72,GREY)
        color.cprint("COUNT [%s] SERVER " % len(result),GREEN)

    def show_pc(self, result):
        # 打印主机配置
        color.cprint("%-5s %-15s %-10s %-10s %-10s %-10s"%("ID","Host","CPU","MEM","Swap","Disk",),YELLOW)
        color.cprint("-"*72,GREY)
        for res in result:
            a = res[0]
            b = res[1]
            c = res[2]
            d = res[3]
            e = res[4]
            f = res[5]
            color.cprint("%-5s %-15s %-10s %-10s %-10s %-10s" % (a,b,c,d,e,f),CYAN)
            #print res,type(res)
        color.cprint("=" * 72,GREY)
        color.cprint("COUNT [%s] PC " % len(result),GREEN)      

    def showplus(self,p):    # 显示所有插件函数
        pp= "show all %s" % p    #("show All %s"%p).upper()
        color.cprint(pp,YELLOW)    #格式化输出
        color.cprint("="*len(pp),GREY)
        if p == 'all':
            sql='select * from pyge'
            self.listpyge(self.sql_all(sql))
        elif p == 'host':
            sql = "select %s from server" % "ip"
            self.show_host(self.sql_all(sql))
        elif p == "multi":
            sql="select * from pyge where type='%s'"%p
            self.listpyge(self.sql_all(sql))
        elif p == "server":
            sql='select id,ip,cpu,memfree,swapfree,diskfree,day,uptime from server'
            self.show_ser(self.sql_all(sql))
        elif p == "pc":
            sql = 'select id,ip,cpu,mem,swap,disk from server'
            self.show_pc(self.sql_all(sql))

    def load(self,plugin):    #载入插件
        def getplu(pid):
            conn=connect(gedb)
            cur=conn.cursor()
            cur.execute('select * from pyge where id=%s'%pid)  
            tmp=cur.fetchone()    # tmp (3, u'multi', u'multi/scan_port_thread')
            cur.close()
            conn.close()
            pat=tmp[2]  # 路径 multi/scan_port_thread 不带后缀
            pty=tmp[1]  # 类型
            if pty == 'payload':
                return ''
            else:
                return pat  # 返回路径
        def noload(p=0):
            if p == 0:
                color.cprint("[!] NO THIS PLUGIN !",RED) 
            else:
                color.cprint("[!] IT'S A PAYLOAD !",RED)
        try:
            pid=int(plugin)   # 要载入的插件的id
            if len(self.sql_all('select * from pyge where id=%s' % pid))==0:  # 如果查询出来的数据为0条
                noload()    # 
            else:
                plu=getplu(pid)   # plu是插件路径  xxx/xxx
                if len(plu)>0:
                    pt=plu.split("/")[0]  # 插件类型
                    load.start(pt,plu)    # 'multi', 'multi/scan_port_thread'
                else:
                    noload(1)
        except:
            if len(self.sql_all('select * from mef where path="%s"'%plugin))==0:
                noload()
            else:
                pt=plugin.split("/")[0]
                load.start(pt,plugin)

    def getplunums(self,p):  # 获取插件数量函数
        if p == 'all':
            return len(self.sql_all('select * from pyge'))  # 所有插件数量
        else:
            return len(self.sql_all('select * from pyge where type="%s"'%p))  # 单一插件数量

    def mainhelp(self):
        color.cprint('PyGE HELP MENU',YELLOW)
        color.cprint('=============',GREY)
        color.cprint('        COMMAND         DESCRIPTION',YELLOW)
        color.cprint('        -------         -----------',GREY,0)
        color.cprint('''
        help            显示帮助信息
        exit            退出  PyGE_App 或者返回上一级
        cls             清空屏幕
        show            显示插件和主机状态
        search          搜索插件
        searchip        搜索主机
        use             使用插件''',CYAN)
        color.cprint('PyGE HELP::SHOW',YELLOW)
        color.cprint('==============',GREY)
        color.cprint('        COMMAND         DESCRIPTION',YELLOW)
        color.cprint('        -------         -----------',GREY,0)
        color.cprint('''
        all             显示又有插件
        host            显示所有主机(仅打印IP)
        server          显示所有服务器的当前状态
        pc              显示所有服务器的基本配置''',CYAN)

    def usage(self,c):
        def ius(c):
            '''def's def =.='''
            color.cprint('[?] USAGE:%s'%c,YELLOW)
        if c == "search":
            ius('search < plugin >')
        if c == "searchip":
            ius('search_host < host >')
        elif c == "show":
            ius('show < all | host | pc | server >')
        elif c == "use":
            ius('use < plugin | pluginID >')

    def banner(self):    # 欢迎画面和插件数量显示函数
        mn=self.getplunums('multi')
        color.cprint('[*] PyGE:: Start Success!!!!!',GREEN)
        #color.cprint('VER::%s::%s'%(nver,seru),CYAN)
        color.cprint('[*] PLU:: Multis::%s' % mn,YELLOW)  # 显示插件数量

    def printge(self):    # PyGE >  显示
        global pyge
        color.cprint(pyge,GREY,0)

    def execmd(self,cmd):  # 执行系统命令
        color.cprint('[*] EXEC:%s'%cmd,RED)
        system(cmd)

    def cls(self):    # 清屏函数
        if name == 'nt':
            system("cls")
        else:
            system("clear")

    def errmsg(self,msg):    # 错误显示函数
        color.cprint("[!] Err:%s"%msg,RED)
        
    def mainexit(self):    # 退出函数
        color.cprint("\n[*] GoodBye :)",RED)
        exit(0)

if __name__=='__main__':
    print __doc__
else:
    cache=cache()
    #cache.start()